# coding : UTF-8
"""
作者：BingBO   时间：2022年08月12日
自动调整代码格式 ：Alt+Ctrl+L
"""
import numpy as np
import cv2 as cv
import glob
import os
import shutil
import time


# 步骤1：采集棋盘格图像（10张以上），并预处理
# 步骤2：找出棋盘格角点坐标
# 步骤3：进一步提取亚像素角点信息
# 步骤4：计算出相机内参数矩阵和畸变系数
# 步骤5：畸变图像校准

def calibrate_camera(checkerboard_image_path=None,
                     checkerboard_row=6,
                     checkerboard_col=9,
                     length=11.5):
    # 世界坐标系中的棋盘格点坐标 , 例如(0,0,0), (1,0,0), (2,0,0) ....,(8,5,0)，去掉Z坐标，记为二维矩阵
    objp = np.zeros((checkerboard_row * checkerboard_col, 3), np.float32)
    objp[:, :2] = np.mgrid[0:checkerboard_col, 0:checkerboard_row].T.reshape(-1, 2)
    objp = objp * length

    # 用于存储所有图像的对象世界点和图像点的数组
    objectPoints = []  # 世界坐标系三维点
    imagePoints = []  # 图像平面二维点

    # 标定所用图像（路径不能有中文）
    # 注：bmp格式无法读取所有
    checkerboard_images = glob.glob(checkerboard_image_path)
    checkerboard_images.sort(key=len)  # 排序

    checkerboard_path_list = checkerboard_images[0::3]  # 棋盘格
    laser1_path_list = checkerboard_images[1::3]  # 棋盘格 + laser1
    laser2_path_list = checkerboard_images[2::3]  # 棋盘格 + laser2
    error_list = []

    found_count = 0  # 记录findCorners success的棋盘格图片数量（没有找到全部内角点的图片不计数）
    found_total = 0
    image_size_col_row = tuple()
    for frame_path in checkerboard_path_list:
        checkerboard_frame = cv.imread(frame_path)
        checkerboard_frame = cv.resize(checkerboard_frame, None, fx=0.4, fy=0.4, interpolation=cv.INTER_CUBIC)
        gray = cv.cvtColor(checkerboard_frame, cv.COLOR_BGR2GRAY)
        image_size_col_row = gray.shape[::-1]  # 倒序(col,row)

        # step1. findChessboardCorners() : 找到棋盘格内角点，存入corners
        found, corners = cv.findChessboardCorners(gray, (checkerboard_col, checkerboard_row), None)

        # 迭代终止条件criteria
        # TERM_CRITERIA_EPS ：满足精度epsilon || TERM_CRITERIA_MAX_ITER ：迭代次数超过阈值max_iter
        criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        found_total += 1
        if found:
            found_count += 1
            objectPoints.append(objp)
            # step2. cornerSubPix() : 提取亚像素角点
            corners_subpix = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
            imagePoints.append(corners_subpix)  # 去掉后面的dtype=float32元素

            # step3. drawChessboardCorners() : 绘制角点 (逐行，每行从左到右)
            cv.drawChessboardCorners(checkerboard_frame, (checkerboard_col, checkerboard_row), corners_subpix, found)
            # cv.imshow('findCorners', checkerboard_frame)
            # cv.waitKey(500)
            print("findCorners success! --> " + "image : " + frame_path)
        else:
            # 将错误的图片\laser1和laser2的错误路径保存
            error_list.append(frame_path)
            error_list.append(laser1_path_list[found_total - 1])  # 索引值记得-1
            error_list.append(laser2_path_list[found_total - 1])
            print("findCorners error! --> " + "image : " + frame_path)
    # 移动错误的图片到error文件夹
    print("error_list:\n", error_list)
    for error_path in error_list:
        filepath, filename = os.path.split(error_path)
        parent_path, _ = os.path.split(filepath)
        newfilepath = os.path.join(parent_path, "error_checkerboard")
        error_move_path = os.path.join(newfilepath, filename)
        shutil.move(error_path, error_move_path)
    print("findCorners Rate:" + str(found_count) + "/" + str(len(checkerboard_path_list)))
    cv.destroyAllWindows()

    # step4. calibrateCamera()
    # 每张图片都有自己的旋转和平移矩阵，但是相机内参和畸变系数只有一组
    ret, cameraMatrix, distCoeffs, R, T = cv.calibrateCamera(objectPoints, imagePoints, image_size_col_row, None, None)
    save_calibrate_result(ret, cameraMatrix, distCoeffs, R, T, imagePoints)

    # for imagePoint in imagePoints:
    #     print("imagePoint:", imagePoint)
    #     print("imagePoint[0]-------------->:", imagePoint[0][0])
    #     print("===================")

    # print("重投影 RMS均方根误差:", ret)
    # print("内参数矩阵:\n", cameraMatrix, '\n')
    # print("畸变系数:\n", distCoeffs, '\n')
    # print("旋转向量(外参数):\n", R, '\n')
    # print("平移向量(外参数):\n", T, '\n')

    mean_error = 0
    total_error = 0
    for i in range(len(objectPoints)):
        # 首先使用cv.projectPoints()将物体点转换为图像点
        # 计算从变换得到的结果和寻角算法之间的绝对范数
        imgPoints2, _ = cv.projectPoints(objectPoints[i], R[i], T[i], cameraMatrix, distCoeffs)
        # imagePoints : 原图中的畸变角点坐标
        # imgpoints2 ：根据R[i], T[i], cameraMatrix, distCoeffs计算的映射角点坐标
        # error = cv.norm(imagePoints[i], imgpoints2, cv.NORM_L2) / len(imgpoints2)
        # mean_error += error  # MSE误差
        error = cv.norm(imagePoints[i], imgPoints2, cv.NORM_L2)
        total_error += error * error
        # 求出平均误差，我们计算了所有标定图像计算误差的算术平均值
    print("RMS error: {}".format(np.sqrt(total_error / (len(objectPoints) * len(imgPoints2)))))
    # print("total error: {}".format(mean_error / len(objectPoints)))


# 得到相机的去畸变图像
def undistort(distortion_image_path):
    distortion_image = cv.imread(distortion_image_path)
    distortion_image = cv.resize(distortion_image, None, fx=0.4, fy=0.4, interpolation=cv.INTER_CUBIC)
    # cv.imshow("distortion_test_image", distortion_test_image)

    cameraMatrix, distCoeffs, R, T, _ = load_calibrate_data()

    h, w = distortion_image.shape[:2]
    # 去畸变前，首先可以根据一个自由缩放参数来细化相机矩阵
    # cv.getOptimalNewCameraMatrix() ：根据给定参数alpha计算最优的新相机内参矩阵，还返回一个图像ROI，可用于裁剪结果

    # 去畸变后的图像四周会出现黑色区域，如果畸变较大，如鱼眼镜头，四周黑色区域将会更大
    # alpha=0，返回具有最小多余像素的未失真图像（去除所有黑色区域这些多余像素）
    # alpha=1，保留所有像素和一些额外的黑色图像
    # alpha=0~1 则得到介于两者之间的效果。

    # cameraMatrix ：原相机的内参矩阵，
    # newCameraMatrix : 畸变图像的相机矩阵，默认情况下与camerammatrix相同，但您可以通过使用不同的矩阵来额外缩放和移动结果。
    # 目标相机（去畸变、无畸变）的内参矩阵，是去畸变后图像的相机的内参矩阵

    # 考虑了畸变系数，
    newCameraMatrix, crop_roi = cv.getOptimalNewCameraMatrix(cameraMatrix, distCoeffs, imageSize=(w, h), alpha=0,
                                                             newImgSize=(w, h))
    print("newCameraMatrix ：\n", newCameraMatrix)

    # step5. undistort() : 畸变图像校准, 与输入畸变图像相同大小（重映射）
    undistort_image = cv.undistort(distortion_image, cameraMatrix, distCoeffs, None, newCameraMatrix)
    cv.imshow("undistort.png", undistort_image)
    cv.imwrite('./result/undistort.png', undistort_image)

    # 方法2
    # mapx, mapy = cv.initUndistortRectifyMap(cameraMatrix, distCoeffs, None, newCameraMatrix, (w, h), 5)
    # undistort_image2 = cv.remap(distortion_image, mapx, mapy, cv.INTER_LINEAR)

    # 裁剪图像
    x, y, w, h = crop_roi
    print("crop_roi ：", crop_roi)
    # 裁去去畸变图像中的 黑色多余像素
    undistort_crop_roi = undistort_image[y:y + h, x:x + w]
    cv.imshow("undistort_crop_roi.png", undistort_crop_roi)
    cv.imwrite('./result/undistort_crop_roi.png', undistort_crop_roi)
    cv.waitKey(0)
    cv.destroyAllWindows()


def save_calibrate_result(ReprojectError, cameraMatrix, distCoeffs, R, T, imagePoints):
    # a：追加，文件存在，文件指针将会放在文件的结尾，新的内容将会被写入到已有内容之后。文件不存在，创建新文件进行写入
    # w+:打开一个文件用于读写，如果该文件已存在则打开文件，并从开头开始编辑，会覆盖原文件的内容；如果该文件不存在，创建新文件
    f = open('./result/标定参数.txt', 'w+')
    f.write("****************************************************\n" +
            "重投影误差(RMS均方根误差):\n" + str(ReprojectError) + '\n' +
            '内参数矩阵:\n' + str(cameraMatrix) + '\n' +
            "畸变系数:\n" + str(distCoeffs) + '\n' +
            "旋转向量(外参数):\n" + str(R) + '\n' +
            "平移向量(外参数):\n" + str(T) + '\n')
    f.close()
    print('内外参、畸变系数已保存 -> ./result/标定参数.txt')

    np.savez('./result/calibrate_data.npz', mtx=cameraMatrix, dist=distCoeffs, rvecs=R, tvecs=T, imagePoints=imagePoints)
    print('calibrate_data.npz文件已保存')


# 从文件加载标定数据
def load_calibrate_data():
    with np.load('./result/calibrate_data.npz') as X:
        cameraMatrix1, distCoeffs1, R, T, imagePoints = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs', 'imagePoints')]
    # print("内参数矩阵:\n", cameraMatrix1, '\n')
    # print("畸变系数:\n", distCoeffs1, '\n')
    # print("旋转向量(外参数):\n", R, '\n')
    # print("平移向量(外参数):\n", T, '\n')
    return cameraMatrix1, distCoeffs1, R, T, imagePoints


# 根据标定结果，映射3D世界坐标系点，到图像平面点
def pose_estimation_demo(checkerboard_image_path=None,
                         checkerboard_row=6,
                         checkerboard_col=9,
                         length=11.5):
    # 世界坐标系中的棋盘格点坐标 , 例如(0,0,0), (1,0,0), (2,0,0) ....,(8,5,0)，去掉Z坐标，记为二维矩阵
    objp = np.zeros((checkerboard_row * checkerboard_col, 3), np.float32)
    objp[:, :2] = np.mgrid[0:checkerboard_col, 0:checkerboard_row].T.reshape(-1, 2)
    objp = objp * length

    # 标定所用图像（路径不能有中文）
    # 注：bmp格式无法读取所有
    checkerboard_images = glob.glob(checkerboard_image_path)
    checkerboard_images.sort(key=len)  # 排序

    cameraMatrix, distCoeffs, R, T, _ = load_calibrate_data()

    found_count = 0  # 记录findCorners success的棋盘格图片数量（没有找到全部内角点的图片不计数）
    found_total = 0
    image_size_col_row = tuple()
    for frame_path in checkerboard_images:
        checkerboard_frame = cv.imread(frame_path)
        checkerboard_frame = cv.resize(checkerboard_frame, None, fx=0.4, fy=0.4, interpolation=cv.INTER_CUBIC)
        gray = cv.cvtColor(checkerboard_frame, cv.COLOR_BGR2GRAY)
        image_size_col_row = gray.shape[::-1]  # 倒序(col,row)

        # step1. findChessboardCorners() : 找到棋盘格内角点，存入corners
        found, corners = cv.findChessboardCorners(gray, (checkerboard_col, checkerboard_row), None)

        # 迭代终止条件criteria
        # TERM_CRITERIA_EPS ：满足精度epsilon || TERM_CRITERIA_MAX_ITER ：迭代次数超过阈值max_iter
        criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        found_total += 1
        if found:
            found_count += 1
            # step2. cornerSubPix() : 提取亚像素角点
            corners_subpix = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)

            # step3. drawChessboardCorners() : 绘制角点 (逐行，每行从左到右)
            cv.drawChessboardCorners(checkerboard_frame, (checkerboard_col, checkerboard_row), corners_subpix, found)
            cv.imshow('findCorners', checkerboard_frame)
            cv.waitKey(500)
            print("findCorners success! --> " + "image : " + frame_path)

            # 计算 旋转向量和平移向量
            ret, rvecs, tvecs = cv.solvePnP(objp, corners_subpix, cameraMatrix, distCoeffs)

            # 映射3D的坐标轴点到图像平面
            # 通过给定的内外参数, 计算三维点投影到二维图像平面上的坐标
            axis = (np.float32([[3, 0, 0], [0, 3, 0], [0, 0, -3]]).reshape(-1, 3)) * length
            # jacobian ：雅可比矩阵
            # imgpts ： 空间点映射后的二维图像点坐标
            imgpts, jac = cv.projectPoints(axis, rvecs, tvecs, cameraMatrix, distCoeffs)
            #
            img = draw(checkerboard_frame, corners_subpix, imgpts)
            cv.imshow('img', img)
            cv.waitKey(0)
        else:
            print("findCorners error! --> " + "image : " + frame_path)
    print("findCorners Rate:" + str(found_count) + "/" + str(len(checkerboard_images)))
    cv.destroyAllWindows()


# 验证2个算法
def test():
    cameraMatrix, distCoeffs, R, T, _ = load_calibrate_data()
    # 世界点
    axis = (np.float32([[3, 0, 0], [0, 3, 0], [0, 0, -3]]).reshape(-1, 3)) * 11.5
    imgpts, jac = cv.projectPoints(axis, R[0], T[0], cameraMatrix, distCoeffs)
    print("（光条上的点）从世界点映射到图像点：", imgpts)
    print(type(imgpts))
    # 对像素坐标去畸变
    undistort_pts = cv.undistortPoints(imgpts, cameraMatrix, distCoeffs, P=cameraMatrix)
    print("去畸变图像点：", undistort_pts)
    # 再用算法 计算图像点映射到世界点的坐标
    print("（光条上的点）从图像坐标映射到点：", imgpts)
    print(type(imgpts))



# 画出三维坐标系
# imgpts ：坐标轴点
def draw(img, corners, imgpts):
    origin = tuple(corners[0].ravel())
    origin = (round(origin[0]), round(origin[1]))
    x_axis = (round(tuple(imgpts[0].ravel())[0]), round(tuple(imgpts[0].ravel())[1]))
    y_axis = (round(tuple(imgpts[1].ravel())[0]), round(tuple(imgpts[1].ravel())[1]))
    z_axis = (round(tuple(imgpts[2].ravel())[0]), round(tuple(imgpts[2].ravel())[1]))
    img = cv.line(img, origin, x_axis, (255, 0, 0), 5)
    img = cv.line(img, origin, y_axis, (0, 255, 0), 5)
    img = cv.line(img, origin, z_axis, (0, 0, 255), 5)

    font = cv.FONT_HERSHEY_SIMPLEX
    cv.putText(img, 'X', (x_axis[0] + 2, x_axis[1] + 2), font, 1, (255, 0, 0), 2, cv.LINE_AA)
    cv.putText(img, 'Y', (y_axis[0] + 2, y_axis[1] + 2), font, 1, (0, 255, 0), 2, cv.LINE_AA)
    cv.putText(img, 'Z', (z_axis[0] + 2, z_axis[1] + 2), font, 1, (0, 0, 255), 2, cv.LINE_AA)

    return img


if __name__ == "__main__":
    # 棋盘格模板规格
    checkerboard_row = 6  # 行 Y (每列点数)
    checkerboard_col = 9  # 列 X (每行点数)
    length = 11.5  # 边长mm
    #
    checkerboard_image_path = r"E:\SHU\Research Group\LaserVisionSensor\Calibrate\checkerboard\*.jpg"
    calibrate_camera(checkerboard_image_path, checkerboard_row, checkerboard_col, length)
    #
    # # 去畸变
    # distortion_image_path = r"E:\SHU\Research Group\LaserVisionSensor\Calibrate\checkerboard\1.jpg"
    # undistort(distortion_image_path)

    # load_calibrate_data()
    # pose_estimation_demo(checkerboard_image_path, checkerboard_row, checkerboard_col, length)

